#include "Head.h"
#include "arcsoft_face_sdk.h"
#include "amcomdef.h"
#include "asvloffscreen.h"
#include "merror.h"

#define USE_V2_1	1
#define USE_RGB		1
#define USE_JPEG	1
#if USE_JPEG
#include "Jpeg.h"
#endif
#include "Tools.h"

#define USE_RAW_FEATURE		0
#define SHOW_DETECT			0
#define SHOW_DETECT_ERROR	0

#define APPID "C4sF7HYmC9nbhc1eb1ubT2rPyGk374b9gwSzq5BA8gDc"
#define SDKKey "FbK1FVZGteDVATj3J9uxfjMJtXUXb8oUWcspf3efn2As"

//nscale值，IMAGE模式下2-32，video模式下2-16，推荐值16
#define NSCALE 16
//检测的人脸数
#define FACENUM	1

#define SafeFree(p) { if ((p)) free(p); (p) = NULL; }
#define SafeArrayDelete(p) { if ((p)) delete [] (p); (p) = NULL; } 
#define SafeDelete(p) { if ((p)) delete (p); (p) = NULL; } 

float thresholdMapping(float confidence)
{
	if((0.0 <= confidence) && (confidence < 0.58)) {
		return confidence*confidence*2.2294887;
	} else if((0.58 <= confidence) && (confidence <= 0.74)) {
		return confidence + 0.17;
	} else if((0.74 < confidence) && (confidence <= 1.0)) {
		return confidence*0.37037037 + 0.62962963;
	} else {
		return 0.0;
	}
}

/**
 * @brief put image names to a vector
 *
 * @param path - path of the image direcotry
 * @param images - the vector of image name
 *
 * @return none
 */
static int ListImages(string const &path, string const &listFileName, vector<string> &images) {
    images.clear();
    // struct dirent *entry;

    /*Check if path is a valid directory path. */
    struct stat s;
    lstat(path.c_str(), &s);
    if (!S_ISDIR(s.st_mode)) {
        fprintf(stderr, "Error: %s is not a valid directory!\n", path.c_str());
        return -1;
    }

    std::ifstream listFile(listFileName.c_str());
    string line;
    while (std::getline(listFile, line))
    {
        images.push_back(line);
    }

    return 0;
}

string baseImagePath = "/media/kasim/Data1/data/ErisedFaceRec/images";
string imageListfile = "/media/kasim/Data1/data/ErisedFaceRec/face_bbox.txt";
string outPath = "/media/kasim/Data1/data/ErisedFaceRec";
bool use_one_out_file = true;

int main(int argc, char** argv)
{
	int threadnum = 12;
	int mode = 0;
    std::cout << "Usage: " << argv[0] << "baseImagePath imageListfile outPath threadnum mode" << std::endl;

    if (argc>1)
        baseImagePath = argv[1];
    if (argc>2)
        imageListfile = argv[2];
    if (argc>3)
        outPath = argv[3];
    if (argc>4)
        threadnum = atoi(argv[4]);
    if (argc>5)
        mode = atoi(argv[5]);

    std::cout << argv[0] << " " << baseImagePath << " " << imageListfile << " " << outPath << " " << threadnum << " " << mode << std::endl;

    vector<string> images;

    /* Load all image names.*/
    ListImages(baseImagePath, imageListfile, images);
    if (images.size() == 0) {
        cerr << "\nError: No images existing under " << baseImagePath << endl;
        return -1;
    }
	int fileCount = images.size();

	if(threadnum<=1)
	{
		threadnum = 1;
		use_one_out_file = true;
	}
    std::string outFileName;
	std::ofstream* outfile = nullptr;
	std::ofstream* outfile2 = nullptr;
	std::ofstream* outerrfile = nullptr;
	if(use_one_out_file)
	{
		outFileName = JoinPathEx(outPath, "Feature.dat");
		outfile = new std::ofstream(outFileName, std::ios::binary);
		if(!outfile->is_open()) {
			cerr << "\nError: Open File " << outFileName << endl;
			return -1;
		}
		outFileName = JoinPathEx(outPath, "BoxFeature.txt");
		outfile2 = new std::ofstream(outFileName, std::ios::out);
		if(!outfile2->is_open()) {
			cerr << "\nError: Open File " << outFileName << endl;
			return -1;
		}
		outFileName = JoinPathEx(outPath, "ExtractError.txt");
		outerrfile = new std::ofstream(outFileName, std::ios::out);
		if(!outerrfile->is_open()) {
			cerr << "\nError: Open File " << outFileName << endl;
			return -1;
		}
	}

	#if USE_V2_1
	MRESULT res = ASFActivation((MPChar)APPID, (MPChar)SDKKey);
	if (MOK != res && MERR_ASF_ALREADY_ACTIVATED != res)
		printf("ASFActivation fail: %ld\n", res);
	else
		printf("ASFActivation sucess: %ld\n", res);
	#else
	MRESULT res = ASFOnlineActivation(APPID, SDKKey);
	if (MOK != res && MERR_ASF_ALREADY_ACTIVATED != res)
		printf("ASFOnlineActivation fail: %d\n", res);
	else
		printf("ASFOnlineActivation sucess: %d\n", res);
	#endif

	auto start_time = GetClockTicks();
	{
		std::mutex m;
		std::condition_variable cv;
		std::thread workers[threadnum];
		static int wake_up_write_thread_id = 0;
		cv.notify_one();
		int frameId = 0;
		for (int i = 0; i < threadnum; i++) {
			workers[i] = thread([&,i]() {
				int idx = i;
				char name[256];

				std::ofstream* _outfile = nullptr;
				std::ofstream* _outfile2 = nullptr;
				std::ofstream* _outerrfile = nullptr;

				if(!use_one_out_file)
				{
					std::string outFileName;
					char _outFileName[128];
					sprintf(_outFileName, "Feature%d.dat", i);
					outFileName = JoinPathEx(outPath, _outFileName);
					_outfile = new std::ofstream(outFileName, std::ios::binary);
					if(!_outfile->is_open()) {
						cerr << "\nError: Open File " << outFileName << endl;
						return;
					}
					sprintf(_outFileName, "BoxFeature%d.txt", i);
					outFileName = JoinPathEx(outPath, _outFileName);
					_outfile2 = new std::ofstream(outFileName, std::ios::out);
					if(!_outfile2->is_open()) {
						cerr << "\nError: Open File " << outFileName << endl;
						return;
					}
					sprintf(_outFileName, "ExtractError%d.txt", i);
					outFileName = JoinPathEx(outPath, _outFileName);
					_outerrfile = new std::ofstream(outFileName, std::ios::out);
					if(!_outerrfile->is_open()) {
						cerr << "\nError: Open File " << outFileName << endl;
						return;
					}
				}
				else
				{
					_outfile = outfile;
					_outfile2= outfile2;
					_outerrfile = outerrfile;
				}

				#if USE_JPEG
				void* jpgDecodeHandle = nullptr;
				jpgDecodeHandle = CreateDecodeJPEG();
				if(nullptr == jpgDecodeHandle)
				{
					return;
				}
				#endif

				//初始化引擎
				MHandle handle = NULL;
				MInt32 mask = ASF_FACERECOGNITION;
				if(mode == 2)
					mask |= ASF_FACE_DETECT;
				MRESULT res = ASFInitEngine(ASF_DETECT_MODE_IMAGE, ASF_OP_0_ONLY, NSCALE, FACENUM, mask, &handle);
				if (res != MOK)
					printf("ASFInitEngine fail: %ld\n", res);
				else
					printf("ASFInitEngine sucess: %ld\n", res);
				int left1, top1, right1, bottom1;

				while (idx < fileCount)
				{
					// if(i == 0)
					// {
					// 	INTERVAL_PROFILE_TIME(ipt0, "extract");
					// }
					ASF_SingleFaceInfo SingleDetectedFaces = { 0 };
					SingleDetectedFaces.faceOrient = 1;
					auto& imageInfo = images[idx];
					if((mode == 1) || (mode == 2))
					{
						sscanf(imageInfo.c_str(), "%s", name);
					}
					else
					{
						sscanf(imageInfo.c_str(), "%s %d,%d,%d,%d\n", name, &left1, &top1, &right1, &bottom1);
					}
					// std::cout << name << std::endl;
					
					std::string imageName = name;
					ASVLOFFSCREEN imageInfo1;
					memset(&imageInfo1, 0, sizeof(ASVLOFFSCREEN));
					#if USE_JPEG
					std::ifstream jpgFile(JoinPathEx(baseImagePath, imageName), std::ios::in | std::ios::binary);
					void* pimage = nullptr;
					int width = 0;
					int height = 0;
					int stride = 0;
					if(jpgFile.is_open())
					{
						std::string buffer((std::istreambuf_iterator<char>(jpgFile)), std::istreambuf_iterator<char>());

						const char* jpgBuf = buffer.c_str();
						int size = buffer.length();
						int org_width = 0;
						int org_height = 0;
						DecodeJPEGYUVEx(jpgDecodeHandle, jpgBuf, size, &pimage, &org_width, &org_height, &stride);
						width = org_width&~0x3;
						height = org_height&~0x1;
						if(right1 >= width)
						{
							right1 = width - 1;
						}
						if(bottom1 >= height)
						{
							bottom1 = height - 1;
						}
						#if !USE_RGB
						if((width != org_width) || (height != org_height))
						{
							const unsigned char* src_y = (const unsigned char*)pimage;
							int uv_stride = stride >> 1;
							const unsigned char* src_u = src_y + height*stride;
							const unsigned char* src_v = src_y + height*stride * 5 / 4;
							
							unsigned char* dst_y = (unsigned char*)malloc(width*height*2);
							int dst_stride = width;
							int dst_uv_stride = dst_stride >> 1;
							unsigned char* dst_u = dst_y + height*dst_stride;
							unsigned char* dst_v = dst_y + height*dst_stride * 5 / 4;
							libyuv::I420Copy(src_y,
											stride,
											src_u,
											uv_stride,
											src_v,
											uv_stride,
											dst_y,
											dst_stride,
											dst_u,
											dst_uv_stride,
											dst_v,
											dst_uv_stride,
											width,
											height);
							stride = width;
							
							DecodeFree(pimage);
							pimage = dst_y;
						}
						#endif
					}
					#if USE_RGB
					int HStride = height;
					int YStride = stride;
					int UStride = YStride>>1;
					int VStride = YStride>>1;
					unsigned char* pYBuf = (unsigned char*)(pimage);
					unsigned char* pUBuf = pYBuf + YStride*HStride;
					unsigned char* pVBuf = pUBuf + UStride*HStride;
					unsigned char* prgbBuf = new unsigned char[YStride*HStride*3];
					int rgbStride = width*3;
					libyuv::I420ToRGB24((const unsigned char*)pYBuf, YStride, (const unsigned char*)pUBuf, UStride,
									(const unsigned char*)pVBuf, VStride, (unsigned char*)prgbBuf, rgbStride,
									width, height);
					imageInfo1.u32PixelArrayFormat = ASVL_PAF_RGB24_B8G8R8;	// 颜色空间格式
					imageInfo1.i32Width		= width;					// 图片宽度
					imageInfo1.i32Height	= height;					// 图片高度
					imageInfo1.ppu8Plane[0]	= prgbBuf;	// 图片数据0
					imageInfo1.pi32Pitch[0]	= rgbStride;					// 图片数据Stride0
					#else
					imageInfo1.u32PixelArrayFormat = ASVL_PAF_I420;	// 颜色空间格式
					imageInfo1.i32Width		= width;					// 图片宽度
					imageInfo1.i32Height	= height;					// 图片高度
					imageInfo1.pi32Pitch[0] = stride;
					imageInfo1.pi32Pitch[1] = stride >> 1;
					imageInfo1.pi32Pitch[2] = stride >> 1;
					imageInfo1.ppu8Plane[0] = (MUInt8*)(pimage);
					imageInfo1.ppu8Plane[1] = imageInfo1.ppu8Plane[0] + imageInfo1.i32Height*stride;
					imageInfo1.ppu8Plane[2] = imageInfo1.ppu8Plane[0] + imageInfo1.i32Height*stride * 5 / 4;
					#endif
					#else
					cv::Mat image1 = cv::imread(JoinPathEx(baseImagePath, imageName));
					Size sz1 = image1.size();
					imageInfo1.u32PixelArrayFormat = ASVL_PAF_RGB24_B8G8R8;	// 颜色空间格式
					imageInfo1.i32Width		= sz1.width;					// 图片宽度
					imageInfo1.i32Height	= sz1.height;					// 图片高度
					imageInfo1.ppu8Plane[0]	= image1.ptr<unsigned char>(0);	// 图片数据0
					imageInfo1.pi32Pitch[0]	= 3*sz1.width;					// 图片数据Stride0
					#endif

					// printf("%s %d,%d,%d,%d\n", name, left1, top1, right1, bottom1);
					bool detect_error = false;
					switch(mode)
					{
					case 2:
						{
							ASF_MultiFaceInfo detectedFaces1 = { 0 };
							#if USE_V2_1
							res = ASFDetectFaces(handle, imageInfo1.i32Width, imageInfo1.i32Height, imageInfo1.u32PixelArrayFormat, (MUInt8*)(imageInfo1.ppu8Plane[0]), &detectedFaces1);
							#else
							res = ASFDetectFacesEx(handle, &imageInfo1, &detectedFaces1);;
							#endif
							if (res != MOK && detectedFaces1.faceNum > 0)
							{
								detect_error = true;
							}
							else
							{
								SingleDetectedFaces.faceRect.left = detectedFaces1.faceRect[0].left;
								SingleDetectedFaces.faceRect.top = detectedFaces1.faceRect[0].top;
								SingleDetectedFaces.faceRect.right = detectedFaces1.faceRect[0].right;
								SingleDetectedFaces.faceRect.bottom = detectedFaces1.faceRect[0].bottom;
								SingleDetectedFaces.faceOrient = detectedFaces1.faceOrient[0];
							}
						}
						break;
					case 1:
                        SingleDetectedFaces.faceRect.left = 0;
                        SingleDetectedFaces.faceRect.top = 0;
                        SingleDetectedFaces.faceRect.right = imageInfo1.i32Width-1;
                        SingleDetectedFaces.faceRect.bottom = imageInfo1.i32Height-1;
					    break;
					case 0:
					default:
                        SingleDetectedFaces.faceRect.left = left1;
                        SingleDetectedFaces.faceRect.top = top1;
                        SingleDetectedFaces.faceRect.right = right1;
                        SingleDetectedFaces.faceRect.bottom = bottom1;
					    break;
					}

					#if SHOW_DETECT && !USE_JPEG
					if(!detect_error)
					{
						cv::rectangle(image1, cv::Point(SingleDetectedFaces.faceRect.left, SingleDetectedFaces.faceRect.top)
										, cv::Point(SingleDetectedFaces.faceRect.right, SingleDetectedFaces.faceRect.bottom), cv::Scalar(0, 255, 0), 2, 1, 0);
						cv::rectangle(image1, cv::Point(left2, top2), cv::Point(right2, bottom2), cv::Scalar(0, 0, 255), 2, 1, 0);
					}
					cv::imshow("image1", image1);
					cv::waitKey();
					#endif

					// 人脸特征提取
					ASF_FaceFeature feature1 = { 0 };

					// 单人脸特征提取
					if(!detect_error)
					{
						// START_PROFILE_TIME(t1, "ASFFaceFeatureExtractEx");
						#if USE_V2_1
						res = ASFFaceFeatureExtract(handle, imageInfo1.i32Width, imageInfo1.i32Height, imageInfo1.u32PixelArrayFormat, (MUInt8*)(imageInfo1.ppu8Plane[0]), &SingleDetectedFaces, &feature1);
						#else
						res = ASFFaceFeatureExtractEx(handle, &imageInfo1, &SingleDetectedFaces, &feature1);
						#endif
						// END_PROFILE_TIME(t1);
					}

					std::unique_lock<std::mutex> lk(m);
					if(use_one_out_file)
					{
						cv.wait(lk, [&,i](){return i==wake_up_write_thread_id;});
					}

					if(detect_error)
					{
						*_outerrfile << imageName << " " << 0 << "," << 0 << "," << 0 << "," << 0 << "," << 0 << "," << res << "\n";
						// _outerrfile->flush();
						#if SHOW_DETECT_ERROR && !USE_JPEG
						cv::imshow("image1", image1);
						cv::waitKey(10);
						#endif
					}
					else
					{
						if (res != MOK) {
							// printf("%s ASFFaceFeatureExtract fail: %ld\n", imageName.c_str(), res);
							*_outerrfile << imageName << " " << SingleDetectedFaces.faceRect.left << "," << SingleDetectedFaces.faceRect.top
										<< "," << SingleDetectedFaces.faceRect.right << "," << SingleDetectedFaces.faceRect.bottom
										<< "," << SingleDetectedFaces.faceOrient << "," << res << "\n";
							// _outerrfile->flush();
							#if SHOW_DETECT_ERROR && !USE_JPEG
							cv::rectangle(image1, cv::Point(SingleDetectedFaces.faceRect.left, SingleDetectedFaces.faceRect.top)
											, cv::Point(SingleDetectedFaces.faceRect.right, SingleDetectedFaces.faceRect.bottom), cv::Scalar(0, 255, 0), 2, 1, 0);
							cv::imshow("image1", image1);
							cv::waitKey(10);
							#endif
						} else {
							// printf("%s ASFFaceFeatureExtract 1 sucess: %ld, featureSize %d\n", imageName.c_str(), res, feature1.featureSize);
							// int featureSize = feature1.featureSize;
							// fwrite(&featureSize, 1, sizeof(int), _outfile);
							#if USE_RAW_FEATURE
							_outfile->write((const char*)(feature1.feature), feature1.featureSize-8);
							*_outfile2 << imageName << " " << SingleDetectedFaces.faceRect.left << "," << SingleDetectedFaces.faceRect.top
										<< "," << SingleDetectedFaces.faceRect.right << "," << SingleDetectedFaces.faceRect.bottom
										<< "," << SingleDetectedFaces.faceOrient << "," << feature1.featureSize-8 << "\n";
							#else
							_outfile->write((const char*)(feature1.feature), feature1.featureSize);
							*_outfile2 << imageName << " " << SingleDetectedFaces.faceRect.left << "," << SingleDetectedFaces.faceRect.top
										<< "," << SingleDetectedFaces.faceRect.right << "," << SingleDetectedFaces.faceRect.bottom
										<< "," << SingleDetectedFaces.faceOrient << "," << feature1.featureSize << "\n";
							#endif
						}
					}

					#if USE_JPEG
					#if USE_RGB
					if(nullptr!=prgbBuf)
					{
						delete[] prgbBuf;
					}
					#endif
					DecodeFree(pimage);
					#endif

					frameId++;
					if(frameId%100 == 0) {
						printf("Process File Count: %d/%d\n", frameId, fileCount);
					}
					if(use_one_out_file)
					{
						wake_up_write_thread_id = i + 1;
						if(wake_up_write_thread_id == threadnum)
						{
							wake_up_write_thread_id = 0;
						}
						lk.unlock();
						// cv.notify_one();
						cv.notify_all();
					}

					idx += threadnum;
				}

				if(!use_one_out_file)
				{
					if(nullptr != _outerrfile)
					{
						_outerrfile->close();
						delete _outerrfile;
						_outerrfile = nullptr;
					}
					if(nullptr != _outfile2)
					{
						_outfile2->close();
						delete _outfile2;
						_outfile2 = nullptr;
					}
					if(nullptr != _outfile)
					{
						_outfile->close();
						delete _outfile;
						_outfile = nullptr;
					}
				}
				#if USE_JPEG
				if(nullptr != jpgDecodeHandle)
				{
					DestroyDecodeJPEG(jpgDecodeHandle);
					jpgDecodeHandle = nullptr;
				}
				#endif
				//销毁引擎
				res = ASFUninitEngine(handle);
				if (res != MOK)
					printf("ASFUninitEngine fail: %ld\n", res);
				else
					printf("ASFUninitEngine sucess: %ld\n", res);
			});
		}

		/* Release thread resources. */
		for (auto &w : workers) {
			if (w.joinable()) w.join();
		}
	}
	auto end_time = GetClockTicks();
	std::cout << "Use Time:" << (float)(end_time - start_time) / 1000.0f << "s" << std::endl;

	if(use_one_out_file)
	{
		if(nullptr != outerrfile)
		{
			outerrfile->close();
			delete outerrfile;
			outerrfile = nullptr;
		}
		if(nullptr != outfile2)
		{
			outfile2->close();
			delete outfile2;
			outfile2 = nullptr;
		}
		if(nullptr != outfile)
		{
			outfile->close();
			delete outfile;
			outfile = nullptr;
		}
	}

	// getchar();
    return 0;
}

